import weka.core.converters.ConverterUtils.DataSource;
import weka.core.Instances;
import weka.core.Instance;
import weka.classifiers.Evaluation;
import weka.classifiers.Classifier;
import weka.classifiers.trees.*;
import weka.classifiers.bayes.*;

import java.io.File;
import java.io.PrintWriter;
import java.io.FileOutputStream;
import java.io.FileInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectInputStream;
import java.util.Arrays;

import java.lang.Class;
import java.lang.reflect.Method;

public class WekaClassifier
{
    public static void main(String[] args) {
        String task = args[0];
        String logfile = args[args.length - 1];
        args = Arrays.copyOfRange(args, 1, args.length - 1);
        WekaClassifier weka = new WekaClassifier();

        if (task.equalsIgnoreCase("build")) {
            weka.build("weka.classifiers.functions.MultilayerPerceptron", args);
        }
        else if (task.equalsIgnoreCase("test")) {
            weka.test(args);
        }
        else if (task.equalsIgnoreCase("classify")) {
            weka.classify(args);
        }
    }
    /*
     * Build a classifier and store it
     * args[0] = train dataset input file
     * args[1] = Classifier object output file
     * args[2..] = options
     */
    protected void build(String algorithm, String[] args) {
        // Train the classifier
        Instances data = loadData(args[0]);

        Object classifier = null;
        try {
            Class <?> clazz = Class.forName(algorithm);
            classifier = clazz.newInstance();
            System.out.println("Init classname");

            Method setOptions = clazz.getDeclaredMethod("setOptions", String[].class);
            String[] options = Arrays.copyOfRange(args, 2, args.length);
            setOptions.invoke(classifier, (Object)options);
            System.out.println("Set options");

            Method buildClassifier = clazz.getDeclaredMethod("buildClassifier", Instances.class);
            buildClassifier.invoke(classifier, data);
            System.out.println("Build classifier");
        }
        catch (Exception ex) {
            System.out.println("Error building classifier, parameters = " + Arrays.toString(args) + ": " + ex.getMessage());
        }

        // Store the classifier
        serializeClassifier(classifier, args[1]);
    }

    /*
     * Restore a classifier and test it
     * args[0] = Train dataset input file
     * args[1] = Classifier object file
     * args[2] = result file
     * args[3..] = options
     */
    protected void test(String[] args) {
        // Restore the classifier
        Classifier classifier = (Classifier) deserializeClassifier(args[1]);

        // Evaluate
        Instances data = loadData(args[0]);
        Evaluation eval = null;
        String test_result = "";
        try {
            eval = new Evaluation(data);
            String[] options = Arrays.copyOfRange(args, 3 - 2, args.length);
            options[0] = "-t";
            options[1] = args[0];
            test_result = eval.evaluateModel(classifier, options);
        }
        catch (Exception ex) {
            System.out.println("Error evaluating: " + ex.getMessage());
        }

        // Output result
        PrintWriter pw = openPrintWriter(args[2]);
        pw.println(test_result);
        pw.close();
    }

    /*
     * Restore a classifier and use it for classification
     * args[0] = test dataset input file
     * args[1] = Classifier object file
     * args[2] = result file
     */
    protected void classify(String[] args) {
        // Restore the classifier
        Classifier classifier = (Classifier) deserializeClassifier(args[1]);

        // Classify
        Instances data = loadData(args[0]);

        PrintWriter pw = openPrintWriter(args[2]);
        for(int i = 0; i < data.numInstances(); ++i) {
            data.instance(i).setClassValue(-1.0);
            double guess = 0.0;
            try {
                guess = classifier.classifyInstance(data.instance(i));
            }
            catch (Exception ex) {
                System.out.println("Error predicting: " + ex.getMessage());
            }

            pw.println((int) guess);
        }
        pw.close();
    }

    /*
     * Loads classification instances from file.
     * Auto detect class / label
     * args = input file name
     */
    private Instances loadData(String args) {
        return loadData(args, -1);
    }
    private Instances loadData(String args, int labelIndex) {
        Instances data = null;
        try {
            DataSource source = new DataSource(args);
            data = source.getDataSet();

            if (labelIndex < 0) { // No label specified. Try to auto detect
                boolean foundLabel = false;
                for(int i = 0; i < data.numAttributes(); ++i)
                    if (data.attribute(i).name().equals("label")
                            && data.attribute(i).isNominal()) {
                        data.setClassIndex(i);
                        foundLabel = true;
                        break;
                    }
                if (!foundLabel) {
                    if (data.attribute(data.numAttributes()-1).isNominal()) {
                        data.setClassIndex(data.numAttributes() - 1);
                    }
                    else {
                        throw new Exception("No class attribute detected");
                    }
                }
            }
            else {
                if (data.attribute(labelIndex).isNominal() == false) {
                    throw new Exception("Invalid class attribute");
                }
            }
        }
        catch (Exception ex) {
            System.out.println("Error reading file " + args + ": " + ex.getMessage());
        }
        return data;
    }

    /*
     * Serialize a Classifier and print it to object file
     * classifier = classifier object
     * filename   = Classifier filename
     */
    private void serializeClassifier(Object classifier, String filename) {
        try {
            ObjectOutputStream oos = new ObjectOutputStream(
                new FileOutputStream(filename)
            );
            oos.writeObject(classifier);
            oos.close();
        }
        catch (Exception ex) {
            System.out.println("Error saving classifier model: " + ex.getMessage());
        }
    }

    /*
     * Deserialize a Classifier and return it
     * filename   = Classifier filename
     */
    private Object deserializeClassifier(String filename) {
        Object classifier = null;
        try {
            ObjectInputStream ois = new ObjectInputStream(
                new FileInputStream(new File(filename))
            );
            classifier = ois.readObject();
            ois.close();
        }
        catch (Exception ex) {
            System.out.println("Error restoring classifier model: " + ex.getMessage());
        }
        return classifier;
    }

    protected PrintWriter openPrintWriter(String filename) {
        PrintWriter pw = null;
        try {
            pw = new PrintWriter(filename);
        }
        catch (Exception ex) {
            System.out.println("Error writing to file: " + filename + " (" + ex.getMessage() + ")");
        }
        return pw;
    }
}
// Build: javac -cp weka.jar WekaClassifier.java
// Build model: java -Xmx2108m -cp weka.jar:. WekaClassifier build output_merged.csv mlp.model build
// Test model:  java -Xmx2108m -cp weka.jar:. WekaClassifier test output_merged.csv mlp.model result.txt test
// Classify:    java -Xmx2108m -cp weka.jar:. WekaClassifier classify output_merged.csv mlp.model result.txt classify
